﻿#!/usr/bin/env python
# -*- coding: utf-8 -*-
#------------------------------------------------------------------------------
# Name:        Distributed "TIE"
# Purpose:     A simple UDP protocol for fast sharing info between node in the
#              same LAN.
#
# Author:      Hu Jun
#
# Created:     27/07/2011
# Copyright:   (c) Hu Jun 2011
# Licence:     GPLv3
# for latest version, visit project's website:
#               http://code.google.com/p/dtie/
#------------------------------------------------------------------------------
import platform
myos=platform.system()

from twisted.internet.protocol import DatagramProtocol
from twisted.internet import reactor
from twisted.application.internet import TimerService
import hashlib
import socket
import time
import struct
from Crypto.Cipher import AES
from Crypto import Random
import chardet
if myos=="Windows":
    import wmi
else:
    import dbus
import wx.lib.newevent


(ResouceUpdateEvt,EVT_RU)=wx.lib.newevent.NewEvent()

def strpad(instr,blocksize=16):
    """Padding input string to multiple of blocksize """
    rstr=instr
    ilen = len(instr)
    mod = ilen % blocksize
    if mod == 0:
        return instr
    else:
        n=((ilen/blocksize) + 1)*blocksize-ilen
        for x in range(n):
            rstr+='_'
        return rstr



class ParseError(Exception):
    def __init__(self,msg,packet):
        self.msg=msg
        self.wrong_packet=packet
    def __str__(self):
        return self.msg

class EncapError(Exception):
    def __init__(self,msg):
        self.msg=msg
    def __str__(self):
        return self.msg

class DTIEProtocol(DatagramProtocol):


    def __init__(self,win,password='DtIE%6&~',lport=50100,grp=1):
        self.password=password
        m=hashlib.md5()
        m.update(self.password)
        m.update(m.digest())
        self.mykey=m.digest()
        self.listening_port=lport
        self.mygroup=grp
        self.packet_seq=0
        self.msg_code_list=[0,1]
        self.resource_type=[0,1,2]
        self.bcast_interval=3
        self.cur_resource=None
        self.max_txt_size=1000
        self.win=win
        self.enableReceving=True

    def prefix2mask(self,prefix):
    	i = 2 ** prefix-1
    	rs=struct.pack('I',i)
    	(c1,c2,c3,c4)=struct.unpack('cccc',rs)
    	rs=''
    	for c in (c1,c2,c3,c4):
    		rs+=str(ord(c))+"."
    	return rs[:-1]

    def getIPList(self):
        global myos
        self.IPList={}
        if myos=="Windows":
            for nic in wmi.WMI ().Win32_NetworkAdapterConfiguration (IPEnabled=1):
                self.IPList[nic.IPAddress[0]]=self.getDBcastAddr(nic.IPAddress[0],
                                                        nic.IPSubnet[0])
            return
        else:
            bus = dbus.SystemBus()
            proxy = bus.get_object("org.freedesktop.NetworkManager","/org/freedesktop/NetworkManager")
            manager = dbus.Interface(proxy, "org.freedesktop.NetworkManager")
            devices = manager.GetDevices()
            for d in devices:
    		   dev_proxy = bus.get_object("org.freedesktop.NetworkManager", d)

    		   prop_iface = dbus.Interface(dev_proxy, "org.freedesktop.DBus.Properties")

    		   # Get the device's current state and interface name
    		   state = prop_iface.Get("org.freedesktop.NetworkManager.Device", "State")
    		   name = prop_iface.Get("org.freedesktop.NetworkManager.Device", "Interface")
    		   ifa = "org.freedesktop.NetworkManager.Device"
    		   type = prop_iface.Get(ifa, "DeviceType")
    		   addr = prop_iface.Get(ifa, "Ip4Address")
    		   v4cfg = prop_iface.Get(ifa, "Ip4Config")
    		   ifb="org.freedesktop.NetworkManager.IP4Config"
    		   v4proxy=bus.get_object("org.freedesktop.NetworkManager", v4cfg)
    		   myiface=dbus.Interface(v4proxy, "org.freedesktop.DBus.Properties")
    		   prefix=myiface.Get(ifb,"Addresses")[0][1]
    		   if state==8:
    			   addr_dot=socket.inet_ntoa(struct.pack('<L',addr))
    			   str_addr=str(addr_dot)
    			   nmask=self.prefix2mask(prefix)
    			   self.IPList[str_addr]=self.getDBcastAddr(str_addr,nmask)

    def getDBcastAddr(self,addr,netmask):
        """return a subnet broadcast address
        """
        alist=addr.split('.')
        nmlist=netmask.split('.')
        if len(alist)<>4 or len(nmlist)<>4:
            return None
##        try:
        a1=int(alist[0])
        if a1==127 or a1>=224 or a1==0:
            return None
        tlist=[]
        for x in alist:
            if int(x)<=0 or int(x)>255:
                return None
            tlist.append(chr(int(x)))
        fms='cccc'
        rs=struct.pack(fms,tlist[0],tlist[1],tlist[2],tlist[3])
        fms='I'
        ai=struct.unpack(fms,rs)
        tlist=[]
        for x in nmlist:
            if int(x)<0 or int(x)>255:
                return None
            tlist.append(chr(int(x)))
        fms='cccc'
        rs=struct.pack(fms,tlist[0],tlist[1],tlist[2],tlist[3])
        fms='I'
        ni=struct.unpack(fms,rs)
##        except:
##            return None


        fms='cccc'
        rs=struct.pack(fms,chr(255),chr(255),chr(255),chr(255))
        fms='I'
        allone=struct.unpack(fms,rs)
        tmpi=allone[0] ^ ni[0]
        dbi=ai[0] | tmpi
        fms='I'
        rs=struct.pack(fms,dbi)
        fms='cccc'
        (s1,s2,s3,s4)=struct.unpack(fms,rs)
        rlist=(s1,s2,s3,s4)
        dbcast=''
        for r in rlist:
            dbcast+=str(ord(r))+'.'
        return dbcast[:-1]



    def setPassword(self,inpass):
        self.password=inpass
        m=hashlib.md5()
        m.update(self.password)
        m.update(m.digest())
        self.mykey=m.digest()

    def setResource(self,rtype,value):
        if not rtype in self.resource_type:
            raise ValueError('rtype must be in range of ',self.resource_type)
            return False
        if value <> None:
            if not isinstance(value,str) and not isinstance(value,unicode):
                raise TypeError('value must be a str')
                return False
            if rtype==0:
                if isinstance(value,str):
                    encoding=chardet.detect(value)['encoding']
                    value=value.decode(encoding)
        self.cur_resource=(rtype,value)

    def startProtocol(self):
        self.transport.socket.setsockopt(socket.SOL_SOCKET,
                                            socket.SO_REUSEADDR, 1)
        self.transport.socket.setsockopt(socket.SOL_SOCKET,
                                            socket.SO_BROADCAST, 1)
        self.getIPList()
        self.bcastRequest()

##        #setting timerservice for broadcast annoucement
##        t1=TimerService(self.bcast_interval,self.bcastAnnounce)
##        t1.startService()

    def sendDatagram(self,v4addr,udpport,packet):
        self.transport.write(packet,(v4addr, udpport))


    def bcastAnnounce(self):
        c_pload=self.gen_attr(self.cur_resource[0],self.cur_resource[1])
        c_pload+=self.gen_attr(3,platform.node())
        e_pload=self.gen_encrypt(c_pload)
        header=self.gen_header(0,len(c_pload))
        packet=self.gen_packet(header,e_pload,c_pload)
        for k,v in self.IPList.items():
            if v<>None:
                self.sendDatagram(v,self.listening_port,packet)

    def sendAnnounce(self,host):
        c_pload=self.gen_attr(self.cur_resource[0],self.cur_resource[1])
        c_pload+=self.gen_attr(3,platform.node())
        e_pload=self.gen_encrypt(c_pload)
        header=self.gen_header(0,len(c_pload))
        packet=self.gen_packet(header,e_pload,c_pload)
        self.sendDatagram(host[0],self.listening_port,packet)

    def bcastRequest(self):
        header=self.gen_header(1,0)
        for k,v in self.IPList.items():
            if v<>None:
                self.sendDatagram(v,self.listening_port,header)


    def setReceving(self,ifcheck):
        self.enableReceving=ifcheck

    def datagramReceived(self, datagram, host):
        """ Executed when a UDP packet is received:
            - datagram: received udp payload; str
            - host: remote host is a tuple of (addr,port); tupple
        """
        if self.enableReceving==False:
            return
        if host[0] in self.IPList.keys(): #drop packet sent from self
            return
        (pheader,a_list)=self.parse_packet(datagram)
        if pheader['code']==1: #request
            if self.cur_resource <> None:
                self.sendAnnounce(host)
            return
        if pheader['code']==0: #announce
            atl=a_list.keys()
            nodename=a_list[3][2]
            atl.remove(3)
            r=a_list[atl[0]]
            evt=ResouceUpdateEvt(node=nodename,Rtype=r[0],Rval=r[2],
                                    IP=host[0],times=time.time())
            wx.PostEvent(self.win,evt)
            return






    def gen_header(self,code,payload_len,group_number=1):
        """ Generate DTIE protocol header:
        - code: message type; unsigned int
        - payload_len: length of encrypted payload, unsigned int
        - group_number
        """
        ver = chr(1) # protocol version is 1
        if not code in range(0,256):
            raise EncapError('code should be in 0-255.')
        mcode = chr(code) # msg type
        if self.packet_seq==65535:
            self.packet_seq=0
        else:
            self.packet_seq+=1
        if not payload_len in range(0,65535):
            raise EncapError('payload len should be in 0-65535.')
        if not group_number in range(0,65535):
            raise EncapError('group_number should be in 0-65535.')
        rheader = struct.pack('ccHHH',ver,mcode,self.packet_seq,group_number,
                                payload_len)
        return rheader

    def gen_attr(self,atype,aval):
        """ Generate dtie protocol attribute:
            - atype: attribute type; int
            - aval: attribute value; str or unicode
        note: aval will be encoded as UTF-8 when atype==0
        """
        if aval <> None:
            if atype==0:
                if isinstance(aval,str):
                    encoding=chardet.detect(aval)['encoding']
                    aval=aval.decode(encoding)
                if isinstance(aval,unicode):
                    aval=aval.encode('utf-8')
            alen=len(aval)
            if (atype==0 or atype==2) and alen>self.max_txt_size:
                fmt_str='HH'
                alen=0
                rattr=struct.pack(fmt_str,atype,alen)
            else:
                fmt_str='HH'+str(alen)+'s'
                rattr=struct.pack(fmt_str,atype,alen,aval)

        else:
            fms='HH'
            rattr=struct.pack(fms,atype,0)
        return rattr

    def gen_encrypt(self,clear_payload):
        """ Generate encrypted payload along with IV
        """

        IV=Random.get_random_bytes(16)
        o=AES.new(self.mykey,AES.MODE_CBC,IV)
        cipher=o.encrypt(strpad(clear_payload))
        return IV+cipher

    def gen_packet(self,header,encrypted_payload='',clear_payload=''):
        m=hashlib.sha256()
        m.update(header+strpad(clear_payload))
        checksum=m.digest()
        return header+encrypted_payload+checksum

    def parse_packet(self,ipacket):
        """ Parse input packet, return:
            - a dict store header info
            - a list store attributes
        """
        plen=len(ipacket)
        if ord(ipacket[1])==1:
            fmt_str='ccHHH'
            (ver,code,seq,group_number,c_plen)=struct.unpack(fmt_str,ipacket)
            if ord(ver)<>1:
                raise ParseError('Unsupported version,packet dropped',ipacket)
                return False
            if c_plen<>0:
                raise ParseError('Invalide packet size,packet dropped',ipacket)
                return False
            if group_number<>self.mygroup:
                raise ParseError('Invalid group number, packet dropped',ipacket)
                return False
            pheader={'ver':ord(ver),'code':ord(code),'seq':seq,
                        'g_num':group_number,'c_plen':c_plen}
            return (pheader,None)

        else:
            if plen<60:
                raise ParseError('Invalide packet size,packet dropped',ipacket)
                return False
            fmt_str='ccHHH16s'+str(plen-56)+'s32s'
            (ver,code,seq,group_number,c_plen,IV,e_pload,checksum)= \
                        struct.unpack(fmt_str,ipacket)
            if ord(ver)<>1:
                raise ParseError('Unsupported version,packet dropped',ipacket)
                return False
            if not ord(code) in self.msg_code_list:
                raise ParseError('Unsupported message code,packet dropped',
                                    ipacket)
                return False
            o=AES.new(self.mykey,AES.MODE_CBC,IV)
            clear_payload=o.decrypt(e_pload)
            fmt_str='8s'+str(plen-8)+'s'
            (header,other)=struct.unpack(fmt_str,ipacket)
            m=hashlib.sha256()
            m.update(header+clear_payload)
            if checksum<>m.digest():
                raise ParseError('Invalide checksum,packet dropped',ipacket)
                return False
            clear_payload=clear_payload[:c_plen] #chop the padding
            alist=self.parse_attrs(clear_payload)
            pheader={'ver':ord(ver),'code':ord(code),'seq':seq,
                        'g_num':group_number,'c_plen':c_plen}

            return (pheader,alist)

    def parse_attrs(self,attrs):
        """ Parse recived attributes, return a dict,
            - key is attribute type
            - value is attribute tuple(TLV)
            Note: returned value will be decode to unicode by using UTF-8
        """
        attr_list={}
        while attrs<>'':
            fmts='HH'
            (atype,alen)=struct.unpack(fmts,attrs[:4])
            fmts='HH'+str(alen)+'s'
            (atype,alen,aval)=struct.unpack(fmts,attrs[:4+alen])
            if atype==0:
                aval=aval.decode('utf-8')
            attr_list[atype]=[atype,alen,aval[:alen]]
            attrs=attrs[4+alen:]

        return attr_list

def main():
    protocol = DTIEProtocol(None)
    protocol.setResource(0,u'这是一个测试！')
    t = reactor.listenUDP(protocol.listening_port, protocol)
    reactor.run()
    print "finish"

if __name__ == '__main__':
    main()